#include <iostream>
#include <fstream>
#include <vector>
#include <chrono>
#include <string>
#include <iomanip>
#include <cstring>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>

// 测试参数结构体
struct TestConfig {
    std::string path;           // 测试文件路径
    size_t file_size_mb;        // 写入的文件大小(MB)
    size_t block_size_kb;       // 写入块大小(KB)
    bool use_direct_io;         // 是否使用O_DIRECT
    bool use_o_sync;            // 是否使用O_SYNC
};

// 单线程写入测试
double test_write_speed(const TestConfig& config) {
    std::string filename = config.path + "/nvme_test.dat";

    int flags = O_WRONLY | O_CREAT | O_TRUNC;
    if (config.use_direct_io) {
        flags |= O_DIRECT;
    }
    if (config.use_o_sync) {
        flags |= O_SYNC;
    }

    int fd = open(filename.c_str(), flags, 0644);
    if (fd < 0) {
        std::cerr << "无法打开文件: " << filename << " 错误: " << strerror(errno) << std::endl;
        return 0.0;
    }

    // 分配4K对齐的内存缓冲区(O_DIRECT需要对齐)
    size_t buffer_size = config.block_size_kb * 1024;
    void* buffer = nullptr;

    if (config.use_direct_io) {
        // 使用posix_memalign分配对齐内存
        if (posix_memalign(&buffer, 4096, buffer_size) != 0) {
            std::cerr << "内存对齐分配失败" << std::endl;
            close(fd);
            return 0.0;
        }
        memset(buffer, 0xff, buffer_size);
    } else {
        // 普通分配
        buffer = malloc(buffer_size);
        memset(buffer, 0xff, buffer_size);
    }

    size_t total_bytes = config.file_size_mb * 1024 * 1024;
    size_t bytes_written = 0;
    size_t write_count = total_bytes / buffer_size;

    auto start_time = std::chrono::high_resolution_clock::now();

    for (size_t i = 0; i < write_count; ++i) {
        ssize_t written = write(fd, buffer, buffer_size);
        if (written < 0) {
            std::cerr << "写入错误: " << strerror(errno) << std::endl;
            break;
        }
        bytes_written += written;
    }

    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();

    // 确保所有数据已写入
    fsync(fd);
    close(fd);
    free(buffer);

    double mb_written = bytes_written / (1024.0 * 1024.0);
    double seconds = duration / 1000.0;
    double mb_per_sec = mb_written / seconds;

    // 删除测试文件
    unlink(filename.c_str());

    return mb_per_sec;
}

// 打印测试结果
void print_results(double speed, const TestConfig& config) {
    std::cout << "========== NVMe写入速率测试结果 ==========" << std::endl;
    std::cout << "测试配置:" << std::endl;
    std::cout << "  - 块大小: " << config.block_size_kb << " KB" << std::endl;
    std::cout << "  - 文件大小: " << config.file_size_mb << " MB" << std::endl;
    std::cout << "  - 使用DirectIO: " << (config.use_direct_io ? "是" : "否") << std::endl;
    std::cout << "  - 使用SyncIO: " << (config.use_o_sync ? "是" : "否") << std::endl;
    std::cout << "\n写入速率: " << std::fixed << std::setprecision(2) << speed << " MB/s" << std::endl;
    std::cout << "==========================================" << std::endl;
}

// 自动确定最佳参数
TestConfig find_best_config(const std::string& path) {
    TestConfig config;
    config.path = path;

    // 测试不同块大小和IO配置
    std::vector<size_t> block_sizes = {64, 128, 256, 512, 1024, 4096, 8192};
    std::vector<bool> direct_io_options = {true, false};
    std::vector<bool> sync_options = {false}; // O_SYNC通常会降低性能

    double best_speed = 0.0;

    std::cout << "正在自动测试最佳参数配置...\n" << std::endl;

    for (auto block_size : block_sizes) {
        for (auto direct_io : direct_io_options) {
            for (auto sync_io : sync_options) {
                TestConfig test_config;
                test_config.path = path;
                test_config.file_size_mb = 256; // 较小的测试样本
                test_config.block_size_kb = block_size;
                test_config.use_direct_io = direct_io;
                test_config.use_o_sync = sync_io;

                std::cout << "测试配置: 块大小=" << block_size << "KB, "
                          << "DirectIO=" << (direct_io ? "是" : "否") << ", "
                          << "SyncIO=" << (sync_io ? "是" : "否") << std::endl;

                double speed = test_write_speed(test_config);
                std::cout << "速度: " << std::fixed << std::setprecision(2) << speed << " MB/s\n" << std::endl;

                if (speed > best_speed) {
                    best_speed = speed;
                    config.block_size_kb = block_size;
                    config.use_direct_io = direct_io;
                    config.use_o_sync = sync_io;
                }
            }
        }
    }

    // 设置其他配置参数
    config.file_size_mb = 1024; // 写入1GB用于最终测试

    std::cout << "找到最佳配置:\n"
              << "  块大小: " << config.block_size_kb << "KB\n"
              << "  DirectIO: " << (config.use_direct_io ? "是" : "否") << "\n"
              << "  SyncIO: " << (config.use_o_sync ? "是" : "否") << "\n"
              << "  测试文件大小: " << config.file_size_mb << "MB\n" << std::endl;

    return config;
}

int main(int argc, char* argv[]) {
    std::cout << "NVMe写入速率测试工具 (单线程版)" << std::endl;

    std::string test_path;
    bool auto_config = true;
    TestConfig config;

    // 解析命令行参数
    if (argc < 2) {
        std::cout << "用法: " << argv[0] << " <测试目录路径> [选项]" << std::endl;
        std::cout << "选项:" << std::endl;
        std::cout << "  --block-size=<KB>      写入块大小(KB)" << std::endl;
        std::cout << "  --file-size=<MB>       文件大小(MB)" << std::endl;
        std::cout << "  --direct=<0|1>         是否使用O_DIRECT" << std::endl;
        std::cout << "  --sync=<0|1>           是否使用O_SYNC" << std::endl;
        std::cout << "  --auto=<0|1>           是否自动寻找最佳参数(默认:1)" << std::endl;
        return 1;
    }

    test_path = argv[1];
    config.path = test_path;

    // 设置默认值
    config.file_size_mb = 1024;
    config.block_size_kb = 1024;
    config.use_direct_io = true;
    config.use_o_sync = false;

    // 解析其他命令行参数
    for (int i = 2; i < argc; i++) {
        std::string arg = argv[i];

        if (arg.find("--block-size=") == 0) {
            config.block_size_kb = std::stoi(arg.substr(13));
            auto_config = false;
        } else if (arg.find("--file-size=") == 0) {
            config.file_size_mb = std::stoi(arg.substr(12));
            auto_config = false;
        } else if (arg.find("--direct=") == 0) {
            config.use_direct_io = (std::stoi(arg.substr(9)) != 0);
            auto_config = false;
        } else if (arg.find("--sync=") == 0) {
            config.use_o_sync = (std::stoi(arg.substr(7)) != 0);
            auto_config = false;
        } else if (arg.find("--auto=") == 0) {
            auto_config = (std::stoi(arg.substr(7)) != 0);
        }
    }

    // 检查目录是否存在并可写
    struct stat st;
    if (stat(test_path.c_str(), &st) != 0 || !S_ISDIR(st.st_mode)) {
        std::cerr << "错误: 指定的路径不是有效目录或不存在: " << test_path << std::endl;
        return 1;
    }

    // 检查是否有写权限
    if (access(test_path.c_str(), W_OK) != 0) {
        std::cerr << "错误: 没有写入权限: " << test_path << std::endl;
        return 1;
    }

    // 自动寻找最佳参数
    if (auto_config) {
        config = find_best_config(test_path);
    }

    std::cout << "开始NVMe写入速率测试...\n" << std::endl;

    // 测试写入性能
    std::cout << "使用最佳配置进行最终测试..." << std::endl;
    double write_speed = test_write_speed(config);

    // 打印结果
    print_results(write_speed, config);

    return 0;
}